"""
基于Transformer的多时间尺度预测策略
策略特点：
1. 同时处理分钟、小时、日线数据
2. 使用多头注意力机制捕捉跨周期关联
3. 实现自适应时间尺度融合
"""

import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import Dataset, DataLoader

class MultiScaleTransformer(nn.Module):
    def __init__(self, input_dims, num_heads=8, hidden_dim=256):
        """
        input_dims: 各时间尺度的输入维度 [minute_dim, hour_dim, day_dim]
        """
        super().__init__()
        self.minute_embed = nn.Linear(input_dims[0], hidden_dim)
        self.hour_embed = nn.Linear(input_dims[1], hidden_dim) 
        self.day_embed = nn.Linear(input_dims[2], hidden_dim)
        
        # 多头注意力层
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads)
        
        # 时间尺度融合层
        self.fusion = nn.Sequential(
            nn.Linear(hidden_dim*3, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)  # 预测收益率
        )
        
    def forward(self, x_minute, x_hour, x_day):
        # 嵌入各时间尺度数据
        minute_feat = self.minute_embed(x_minute)
        hour_feat = self.hour_embed(x_hour)
        day_feat = self.day_embed(x_day)
        
        # 跨时间尺度注意力
        all_feat = torch.cat([minute_feat.unsqueeze(0), 
                             hour_feat.unsqueeze(0),
                             day_feat.unsqueeze(0)], dim=0)
        
        attn_out, _ = self.attention(all_feat, all_feat, all_feat)
        
        # 融合特征
        fused = torch.cat([attn_out[0], attn_out[1], attn_out[2]], dim=1)
        return self.fusion(fused)

class MultiScaleDataset(Dataset):
    def __init__(self, minute_data, hour_data, day_data, labels):
        self.minute_data = torch.FloatTensor(minute_data)
        self.hour_data = torch.FloatTensor(hour_data)
        self.day_data = torch.FloatTensor(day_data)
        self.labels = torch.FloatTensor(labels)
        
    def __len__(self):
        return len(self.labels)
        
    def __getitem__(self, idx):
        return (self.minute_data[idx], 
                self.hour_data[idx],
                self.day_data[idx],
                self.labels[idx])

class MultiScaleStrategy:
    def __init__(self, asset_name, lookback=20):
        self.asset = asset_name
        self.model = MultiScaleTransformer([10, 8, 5])  # 示例维度
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
        self.criterion = nn.MSELoss()
        
    def train(self, dataloader, epochs=50):
        self.model.train()
        for epoch in range(epochs):
            for x_min, x_hr, x_day, y in dataloader:
                self.optimizer.zero_grad()
                pred = self.model(x_min, x_hr, x_day)
                loss = self.criterion(pred.squeeze(), y)
                loss.backward()
                self.optimizer.step()
                
    def predict(self, minute_data, hour_data, day_data):
        self.model.eval()
        with torch.no_grad():
            pred = self.model(
                torch.FloatTensor(minute_data),
                torch.FloatTensor(hour_data), 
                torch.FloatTensor(day_data))
        return pred.item()
        
    def generate_signals(self, pred, threshold=0.005):
        """根据预测生成交易信号"""
        if pred > threshold:
            return 1  # 买入
        elif pred < -threshold:
            return -1  # 卖出
        return 0  # 持有

if __name__ == '__main__':
    # 示例用法
    np.random.seed(42)
    
    # 模拟多时间尺度数据
    minute_data = np.random.randn(1000, 10)
    hour_data = np.random.randn(1000, 8)
    day_data = np.random.randn(1000, 5)
    labels = np.random.randn(1000) * 0.01  # 模拟收益率
    
    # 创建数据集
    dataset = MultiScaleDataset(minute_data, hour_data, day_data, labels)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
    
    # 训练模型
    strategy = MultiScaleStrategy("AAPL")
    strategy.train(dataloader)
    
    # 测试预测
    test_min = np.random.randn(10)
    test_hr = np.random.randn(8) 
    test_day = np.random.randn(5)
    pred = strategy.predict(test_min, test_hr, test_day)
    print(f"预测收益率: {pred:.4f}")
    print("交易信号:", strategy.generate_signals(pred))